package net.cyue.ort.llm.util;

import java.nio.ByteBuffer;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

/**
 * 资源池
 * 管理ByteBuffer的复用，减少内存分配和回收
 */
public class ResourcePool {
    private final BlockingQueue<ByteBuffer> bufferPool = new ArrayBlockingQueue<>(100);
    private final int maxBufferSize;
    
    public ResourcePool(int maxBufferSize) {
        this.maxBufferSize = maxBufferSize;
    }
    
    /**
     * 借用一个缓冲区
     * @param size 需要的缓冲区大小
     * @return 缓冲区
     */
    public ByteBuffer borrowBuffer(int size) {
        if (size > maxBufferSize) {
            throw new IllegalArgumentException("Buffer size too large: " + size);
        }
        
        ByteBuffer buffer = bufferPool.poll();
        if (buffer == null || buffer.capacity() < size) {
            return ByteBuffer.allocateDirect(size);
        }
        
        buffer.clear();
        buffer.limit(size);
        return buffer;
    }
    
    /**
     * 归还缓冲区到池中
     * @param buffer 要归还的缓冲区
     */
    public void returnBuffer(ByteBuffer buffer) {
        if (buffer == null) {
            return;
        }
        
        if (buffer.capacity() <= maxBufferSize && bufferPool.size() < 100) {
            buffer.clear();
            bufferPool.offer(buffer);
        }
    }
    
    /**
     * 清理池中的所有缓冲区
     */
    public void clear() {
        bufferPool.clear();
    }
    
    /**
     * 获取当前池中缓冲区数量
     */
    public int size() {
        return bufferPool.size();
    }
}

